#! /usr/bin/env python
#
# couchdb-xapian-indexer - Index couchdb databases with Xapian.
# Copyright (C) 2008 Paul J. Davis <paul.joseph.davis@gmail.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import logging
import os
import sys
from optparse import OptionParser, make_option

import couchdb
import simplejson
import xapian

log = logging.getLogger(__name__)

class Index(object):
    def __init__(self, dir, url, dbname):
        log.info("Creating indexer for: %s" % dbname)
        self.idxfile = os.path.join(dir, '%s.idx' % dbname)
        self.server = couchdb.Server(url)
        self.dbname = dbname
        self.xdb = xapian.WritableDatabase(self.idxfile, xapian.DB_CREATE_OR_OPEN)
        self.index = xapian.TermGenerator()
        self.index.set_stemmer(xapian.Stem("english"))
        self.startkey = self.xdb.get_metadata("startkey")
        if self.startkey:
            self.startkey = simplejson.loads(self.startkey)
        self.batch_size = 1000
    def reindex(self):
        log.info("Reindexing %s" % self.dbname)
        try:
            if self.dbname not in self.server:
                return
            db = self.server[self.dbname]
            self.xdb.begin_transaction()
            params = {'count': self.batch_size}
            if self.startkey:
                params['startkey'] = self.startkey
            docs = db.view('_all_docs_by_seq', **params)
            while len(docs) > 0:
                for doc in docs:
                    self.startkey = doc.key
                    xuid = "COUCHDB_ID_%s" % doc.id
                    if doc.value.get('deleted', False):
                        self.xdb.delete_document(doc.id)
                    else:
                        cdoc = db[doc.id]
                        xdoc = xapian.Document()
                        xdoc.add_term(xuid)
                        xdoc.set_data(doc.id)
                        self.index.set_document(xdoc)
                        self.index.index_text(' '.join(map(lambda x: str(cdoc.get(x, '')), [k for k in cdoc])))
                        self.xdb.replace_document(xuid, xdoc)
                params['startkey'] = self.startkey
                docs = db.view('_all_docs_by_seq', **params)
            self.xdb.set_metadata("startkey", simplejson.dumps(self.startkey))
            self.xdb.commit_transaction()
        except:
            self.xdb.cancel_transaction()
            raise

def updates():
    line = sys.stdin.readline()
    while line:
        if not line:
            return
        obj = simplejson.loads(line)
        yield obj
        line = sys.stdin.readline()

def main(dir, url, exclude):
    if not os.path.isdir(dir):
        os.mkdir(dir)
    indices = {}
    for update in updates():
        dbname = update['db']
        if dbname in exclude:
            continue
        if dbname not in indices:
            indices[dbname] = Index(dir, url, dbname)
        indices[dbname].reindex()

if __name__ == '__main__':
    options = [
        make_option('-d', '--dir', dest='dir', metavar="DIRECTORY", default="./xapian",
            help="Directory in which to store xapian databases. [%default]"),
        make_option('-u', '--url', dest='url', metavar="URL", default="http://localhost:5984",
            help="URL of the couchdb server. [%default]"),
        make_option('-e', '--exclude', dest='exclude', metavar='DB_NAME', default=[],
            help="Exclude a database from indexing. Can be used multiple times."),
        make_option('-l', '--log', dest='log', metavar="FILE", default='./xapian/index.log',
            help="Name of the log file to write to."),
    ]
    parser = OptionParser("usage: %prog [OPTIONS]", option_list=options)
    opts, args = parser.parse_args()
    if len(args) != 0:
        print "Unrecognized arguments: %s" % ' '.join(args)
        parser.print_help()
        exit(-1)
    logging.basicConfig(filename=opts.log, level=logging.DEBUG, format="%(levelname)s %(message)s")
    try:
        main(os.path.abspath(opts.dir), opts.url, opts.exclude)
    except:
        log.exception("Indexer shutting down due to high stress. Relaxation needed.")
